#!/usr/bin/env python3
"""
CogView MCP Server

A Model Context Protocol server that provides image generation using CogView
with optional background removal functionality.
"""

import asyncio
import os
import io
import tempfile
from typing import Any, Sequence, Optional
from contextlib import asynccontextmanager
import logging
import json
import sys

try:
    from mcp.server.models import InitializationOptions
    from mcp.server import NotificationOptions, Server
    from mcp.types import (
        Resource,
        Tool,
        TextContent,
        ImageContent,
        EmbeddedResource,
    )
    from pydantic import AnyUrl
    import mcp.types as types
    from mcp.server.stdio import stdio_server
    MCP_AVAILABLE = True
except ImportError:
    # Fallback for MCP import issues
    MCP_AVAILABLE = False

from .cogview_client import CogViewClient
from .background_remover import BackgroundRemover
from .config import Config

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class CogViewMCPServer:
    """MCP Server for CogView image generation with background removal."""
    
    def __init__(self):
        if not MCP_AVAILABLE:
            logger.warning("MCP libraries not available. Running in standalone mode.")
            self.cogview_client = CogViewClient()
            self.bg_remover = BackgroundRemover()
            return
            
        self.server = Server("cogview-mcp-server")
        self.cogview_client = CogViewClient()
        self.bg_remover = BackgroundRemover()
        
    async def initialize(self):
        """Initialize the server and register handlers."""
        
        @self.server.list_tools()
        async def handle_list_tools() -> list[Tool]:
            """List available tools."""
            return [
                Tool(
                    name="generate_image",
                    description="Generate an image using CogView based on a text prompt",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "prompt": {
                                "type": "string",
                                "description": "Text prompt describing the image to generate"
                            },
                            "model": {
                                "type": "string",
                                "description": "CogView model to use",
                                "default": "cogview-3-flash",
                                "enum": ["cogview-3-flash"]
                            },
                            "size": {
                                "type": "string",
                                "description": "Image size in format 'widthxheight'",
                                "default": "1024x1024",
                                "enum": ["1024x1024", "768x1024", "1024x768", "1152x896", "896x1152", "1216x832", "832x1216", "1344x768", "768x1344", "1536x640", "640x1536"]
                            }
                        },
                        "required": ["prompt"]
                    }
                ),
                Tool(
                    name="generate_image_transparent",
                    description="Generate an image with transparent background by removing the background from a CogView generated image",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "prompt": {
                                "type": "string",
                                "description": "Text prompt describing the image to generate"
                            },
                            "model": {
                                "type": "string",
                                "description": "CogView model to use",
                                "default": "cogview-3-flash",
                                "enum": ["cogview-3-flash"]
                            },
                            "size": {
                                "type": "string",
                                "description": "Image size in format 'widthxheight'",
                                "default": "1024x1024",
                                "enum": ["1024x1024", "768x1024", "1024x768", "1152x896", "896x1152", "1216x832", "832x1216", "1344x768", "768x1344", "1536x640", "640x1536"]
                            },
                            "background_removal_model": {
                                "type": "string",
                                "description": "Background removal model to use",
                                "default": "u2net",
                                "enum": ["u2net", "u2netp", "u2net_human_seg", "silueta"]
                            }
                        },
                        "required": ["prompt"]
                    }
                )
            ]

        @self.server.call_tool()
        async def handle_call_tool(name: str, arguments: dict[str, Any] | None) -> list:
            """Handle tool calls."""
            if arguments is None:
                arguments = {}

            try:
                if name == "generate_image":
                    return await self._generate_image(arguments)
                elif name == "generate_image_transparent":
                    return await self._generate_image_transparent(arguments)
                else:
                    raise ValueError(f"Unknown tool: {name}")
                    
            except Exception as e:
                logger.error(f"Error in tool {name}: {str(e)}")
                if MCP_AVAILABLE:
                    return [types.TextContent(
                        type="text",
                        text=f"Error generating image: {str(e)}"
                    )]
                else:
                    return [{"type": "text", "text": f"Error generating image: {str(e)}"}]

    async def _generate_image(self, arguments: dict[str, Any]) -> list:
        """Generate an image using CogView."""
        prompt = arguments.get("prompt", "")
        model = arguments.get("model", "cogview-3-flash")
        size = arguments.get("size", Config.DEFAULT_IMAGE_SIZE)
        
        if not prompt:
            if MCP_AVAILABLE:
                return [types.TextContent(
                    type="text",
                    text="Error: prompt is required"
                )]
            else:
                return [{"type": "text", "text": "Error: prompt is required"}]
        
        logger.info(f"Generating image with prompt: {prompt}, model: {model}, size: {size}")
        
        try:
            # Generate image using CogView
            image_url = await self.cogview_client.generate_image(prompt, model, size)
            
            # Download and encode image
            image_data = await self.cogview_client.download_image(image_url)
            
            if MCP_AVAILABLE:
                return [
                    types.TextContent(
                        type="text",
                        text=f"Image generated successfully with prompt: '{prompt}'"
                    ),
                    types.ImageContent(
                        type="image",
                        data=image_data,
                        mimeType="image/png"
                    )
                ]
            else:
                return [
                    {"type": "text", "text": f"Image generated successfully with prompt: '{prompt}'"},
                    {"type": "image", "data": image_data, "mimeType": "image/png"}
                ]
            
        except Exception as e:
            logger.error(f"Error generating image: {str(e)}")
            if MCP_AVAILABLE:
                return [types.TextContent(
                    type="text",
                    text=f"Error generating image: {str(e)}"
                )]
            else:
                return [{"type": "text", "text": f"Error generating image: {str(e)}"}]

    async def _generate_image_transparent(self, arguments: dict[str, Any]) -> list:
        """Generate an image with transparent background."""
        prompt = arguments.get("prompt", "")
        model = arguments.get("model", "cogview-3-flash")
        size = arguments.get("size", Config.DEFAULT_IMAGE_SIZE)
        bg_model = arguments.get("background_removal_model", "u2net")
        
        if not prompt:
            if MCP_AVAILABLE:
                return [types.TextContent(
                    type="text",
                    text="Error: prompt is required"
                )]
            else:
                return [{"type": "text", "text": "Error: prompt is required"}]
        
        logger.info(f"Generating transparent image with prompt: {prompt}, model: {model}, size: {size}")
        
        try:
            # Generate image using CogView
            image_url = await self.cogview_client.generate_image(prompt, model, size)
            
            # Download image
            image_data = await self.cogview_client.download_image(image_url)
            
            # Remove background
            transparent_image_data = await self.bg_remover.remove_background(image_data, bg_model)
            
            if MCP_AVAILABLE:
                return [
                    types.TextContent(
                        type="text",
                        text=f"Transparent image generated successfully with prompt: '{prompt}'"
                    ),
                    types.ImageContent(
                        type="image",
                        data=transparent_image_data,
                        mimeType="image/png"
                    )
                ]
            else:
                return [
                    {"type": "text", "text": f"Transparent image generated successfully with prompt: '{prompt}'"},
                    {"type": "image", "data": transparent_image_data, "mimeType": "image/png"}
                ]
            
        except Exception as e:
            logger.error(f"Error generating transparent image: {str(e)}")
            if MCP_AVAILABLE:
                return [types.TextContent(
                    type="text",
                    text=f"Error generating transparent image: {str(e)}"
                )]
            else:
                return [{"type": "text", "text": f"Error generating transparent image: {str(e)}"}]

    async def run(self):
        """Run the MCP server."""
        if not MCP_AVAILABLE:
            logger.error("MCP libraries not available. Please install with: pip install mcp")
            return
            
        await self.initialize()
        
        # Import and set up the transport
        async with stdio_server() as (read_stream, write_stream):
            await self.server.run(
                read_stream,
                write_stream,
                InitializationOptions(
                    server_name=Config.SERVER_NAME,
                    server_version=Config.SERVER_VERSION,
                    capabilities=self.server.get_capabilities(
                        notification_options=NotificationOptions(),
                        experimental_capabilities={}
                    )
                )
            )


async def main():
    """Main entry point."""
    server = CogViewMCPServer()
    await server.run()


if __name__ == "__main__":
    asyncio.run(main())
